package cn.harry.modular.ai.service.impl;

import cn.harry.common.api.R;
import cn.harry.common.exception.ApiException;
import cn.harry.common.utils.BeanUtils;
import cn.harry.modular.ai.core.enums.AiPlatformEnum;
import cn.harry.modular.ai.core.util.AiUtils;
import cn.harry.modular.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
import cn.harry.modular.ai.domain.ChatConversation;
import cn.harry.modular.ai.domain.ChatMessage;
import cn.harry.modular.ai.enums.AiChatRoleEnum;
import cn.harry.modular.ai.enums.ErrorCodeConstants;
import cn.harry.modular.ai.mapper.ChatMessageMapper;
import cn.harry.modular.ai.service.ApiKeyService;
import cn.harry.modular.ai.service.ChatConversationService;
import cn.harry.modular.ai.service.ChatMessageService;
import cn.harry.modular.ai.service.ChatModelService;
import cn.harry.modular.ai.vo.message.AiChatMessageSendReqVO;
import cn.harry.modular.ai.vo.message.AiChatMessageSendRespVO;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjUtil;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import jakarta.annotation.Resource;
import jakarta.validation.constraints.NotEmpty;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

import static cn.harry.common.utils.CollectionUtils.convertList;
import static cn.harry.modular.ai.enums.ErrorCodeConstants.CHAT_CONVERSATION_NOT_EXISTS;

/**
 * AI 聊天消息服务实现类
 *
 * @author harry
 * @公众号 Harry技术
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class ChatMessageServiceImpl extends ServiceImpl<ChatMessageMapper, ChatMessage> implements ChatMessageService {

    @Resource
    private ChatMessageMapper chatMessageMapper;
    @Resource
    private ChatConversationService chatConversationService;
    @Resource
    private ChatModelService chatModalService;
    @Resource
    private ApiKeyService apiKeyService;

    @Override
    public AiChatMessageSendRespVO sendMessage(AiChatMessageSendReqVO sendReqVO, Long userId) {
        // 1.1 校验对话存在
        ChatConversation conversation = chatConversationService.validateChatConversationExists(sendReqVO.getConversationId());
//        if (ObjUtil.notEqual(conversation.getUserId(), userId)) {
//            throw new ApiException(CHAT_CONVERSATION_NOT_EXISTS);
//        }
        List<ChatMessage> historyMessages = list(Wrappers.<ChatMessage>lambdaQuery().eq(ChatMessage::getConversationId, conversation.getId()));
        // 1.2 校验模型

        cn.harry.modular.ai.domain.ChatModel model = chatModalService.validateChatModel(conversation.getModelId());
        ChatModel chatModel = apiKeyService.getChatModel(model.getKeyId());

        // 2. 插入 user 发送消息
        ChatMessage userMessage = createChatMessage(conversation.getId(), null, model,
                userId, conversation.getRoleId(), MessageType.USER, sendReqVO.getContent(), sendReqVO.getUseContext());

        // 3.1 插入 assistant 接收消息
        ChatMessage assistantMessage = createChatMessage(conversation.getId(), userMessage.getId(), model,
                userId, conversation.getRoleId(), MessageType.ASSISTANT, "", sendReqVO.getUseContext());

        // 3.2 召回段落 TODO 暂时不召回
        List<AiKnowledgeSegmentDO> segmentList = Collections.emptyList();

        // 3.3 创建 chat 需要的 Prompt
        Prompt prompt = buildPrompt(conversation, historyMessages, segmentList, model, sendReqVO);
        ChatResponse chatResponse = chatModel.call(prompt);

        // 3.4 段式返回
        String newContent = chatResponse.getResult().getOutput().getContent();

        ChatMessage msg = new ChatMessage();
        msg.setId(assistantMessage.getId());
        msg.setContent(newContent);
        msg.setSegmentIds(convertList(segmentList, AiKnowledgeSegmentDO::getId));
        updateById(msg);

        AiChatMessageSendRespVO respVO = new AiChatMessageSendRespVO();
        respVO.setSend(BeanUtils.toBean(userMessage, AiChatMessageSendRespVO.Message.class));
        AiChatMessageSendRespVO.Message receive = BeanUtils.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class);
        receive.setContent(newContent);
        respVO.setReceive(receive);
        return respVO;
    }

    @Override
    public Flux<R<AiChatMessageSendRespVO>> sendChatMessageStream(AiChatMessageSendReqVO sendReqVO, Long userId) {
        // 1.1 校验对话存在
        ChatConversation conversation = chatConversationService.validateChatConversationExists(sendReqVO.getConversationId());
        if (ObjUtil.notEqual(conversation.getUserId(), userId)) {
            throw new ApiException(CHAT_CONVERSATION_NOT_EXISTS);
        }
        List<ChatMessage> historyMessages = list(Wrappers.<ChatMessage>lambdaQuery().eq(ChatMessage::getConversationId, conversation.getId()));
        // 1.2 校验模型

        cn.harry.modular.ai.domain.ChatModel model = chatModalService.validateChatModel(conversation.getModelId());
        ChatModel chatModel = apiKeyService.getChatModel(model.getKeyId());

        // 2. 插入 user 发送消息
        ChatMessage userMessage = createChatMessage(conversation.getId(), null, model,
                userId, conversation.getRoleId(), MessageType.USER, sendReqVO.getContent(), sendReqVO.getUseContext());

        // 3.1 插入 assistant 接收消息
        ChatMessage assistantMessage = createChatMessage(conversation.getId(), userMessage.getId(), model,
                userId, conversation.getRoleId(), MessageType.ASSISTANT, "", sendReqVO.getUseContext());

        // 3.2 召回段落 TODO 暂时不召回
        List<AiKnowledgeSegmentDO> segmentList = Collections.emptyList();

        // 3.3 创建 chat 需要的 Prompt
        Prompt prompt = buildPrompt(conversation, historyMessages, segmentList, model, sendReqVO);
        Flux<ChatResponse> streamResponse = chatModel.stream(prompt);


        // 3.4 流式返回
        // TODO 注意：Schedulers.immediate() 目的是，避免默认 Schedulers.parallel() 并发消费 chunk 导致 SSE 响应前端会乱序问题
        StringBuffer contentBuffer = new StringBuffer();
        return streamResponse.map(chunk -> {
            String newContent = chunk.getResult() != null ? chunk.getResult().getOutput().getContent() : null;
            newContent = StrUtil.nullToDefault(newContent, ""); // 避免 null 的 情况
            contentBuffer.append(newContent);
            // 响应结果

            AiChatMessageSendRespVO respVO = new AiChatMessageSendRespVO();
            respVO.setSend(BeanUtils.toBean(userMessage, AiChatMessageSendRespVO.Message.class));
            AiChatMessageSendRespVO.Message receive = BeanUtils.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class);
            receive.setContent(newContent);
            respVO.setReceive(receive);

            return R.success(respVO);

        }).doOnComplete(() -> {


            ChatMessage msg = new ChatMessage();
            msg.setId(assistantMessage.getId());
            msg.setContent(contentBuffer.toString());
            msg.setSegmentIds(convertList(segmentList, AiKnowledgeSegmentDO::getId));
            updateById(msg);

        }).doOnError(throwable -> {
            log.error("[sendChatMessageStream][userId({}) sendReqVO({}) 发生异常]", userId, sendReqVO, throwable);


            ChatMessage msg = new ChatMessage();
            msg.setId(assistantMessage.getId());
            msg.setContent(throwable.getMessage());
            updateById(msg);
        }).onErrorResume(error -> Flux.just(
                R.failed(ErrorCodeConstants.CHAT_STREAM_ERROR)));
    }

    @Override
    public List<ChatMessage> getChatMessageListByConversationId(Long conversationId) {
        return list(Wrappers.<ChatMessage>lambdaQuery()
                .eq(ChatMessage::getConversationId, conversationId)
                .orderByAsc(ChatMessage::getId)
        );
    }


    private Prompt buildPrompt(ChatConversation conversation,
                               List<ChatMessage> messages,
                               List<AiKnowledgeSegmentDO> segmentList,
                               cn.harry.modular.ai.domain.ChatModel model,
                               AiChatMessageSendReqVO sendReqVO) {
        // 1. 构建 Prompt Message 列表
        List<Message> chatMessages = new ArrayList<>();

        // 1.1 召回内容消息构建
        if (CollUtil.isNotEmpty(segmentList)) {
            PromptTemplate promptTemplate = new PromptTemplate(AiChatRoleEnum.AI_KNOWLEDGE_ROLE.getSystemMessage());
            StringBuilder infoBuilder = StrUtil.builder();
            segmentList.forEach(segment -> infoBuilder.append(System.lineSeparator()).append(segment.getContent()));
            Message message = promptTemplate.createMessage(Map.of("info", infoBuilder.toString()));
            chatMessages.add(message);
        }

        // 1.2 system context 角色设定
        if (StrUtil.isNotBlank(conversation.getSystemMessage())) {
            chatMessages.add(new SystemMessage(conversation.getSystemMessage()));
        }
        // 1.3 history message 历史消息
        List<ChatMessage> contextMessages = filterContextMessages(messages, conversation, sendReqVO);
        contextMessages.forEach(message -> chatMessages.add(AiUtils.buildMessage(message.getType(), message.getContent())));
        // 1.4 user message 新发送消息
        chatMessages.add(new UserMessage(sendReqVO.getContent()));

        // 2. 构建 ChatOptions 对象
        AiPlatformEnum platform = AiPlatformEnum.validatePlatform(model.getPlatform());
        ChatOptions chatOptions = AiUtils.buildChatOptions(platform, model.getModel(),
                conversation.getTemperature(), conversation.getMaxTokens());
        return new Prompt(chatMessages, chatOptions);
    }

    /**
     * 从历史消息中，获得倒序的 n 组消息作为消息上下文
     * <p>
     * n 组：指的是 user + assistant 形成一组
     *
     * @param messages     消息列表
     * @param conversation 对话
     * @param sendReqVO    发送请求
     * @return 消息上下文
     */
    private List<ChatMessage> filterContextMessages(List<ChatMessage> messages, ChatConversation conversation, AiChatMessageSendReqVO sendReqVO) {
        if (conversation.getMaxContexts() == null || ObjUtil.notEqual(sendReqVO.getUseContext(), Boolean.TRUE)) {
            return Collections.emptyList();
        }
        List<ChatMessage> contextMessages = new ArrayList<>(conversation.getMaxContexts() * 2);
        for (int i = messages.size() - 1; i >= 0; i--) {
            ChatMessage assistantMessage = CollUtil.get(messages, i);
            if (assistantMessage == null || assistantMessage.getReplyId() == null) {
                continue;
            }
            ChatMessage userMessage = CollUtil.get(messages, i - 1);
            if (userMessage == null || ObjUtil.notEqual(assistantMessage.getReplyId(), userMessage.getId())
                    || StrUtil.isEmpty(assistantMessage.getContent())) {
                continue;
            }
            // 由于后续要 reverse 反转，所以先添加 assistantMessage
            contextMessages.add(assistantMessage);
            contextMessages.add(userMessage);
            // 超过最大上下文，结束
            if (contextMessages.size() >= conversation.getMaxContexts() * 2) {
                break;
            }
        }
        Collections.reverse(contextMessages);
        return contextMessages;
    }

    private ChatMessage createChatMessage(Long conversationId,
                                          Long replyId,
                                          cn.harry.modular.ai.domain.ChatModel model,
                                          Long userId, Long roleId,
                                          MessageType messageType,
                                          @NotEmpty(message = "聊天内容不能为空") String content,
                                          Boolean useContext) {
        ChatMessage chatMessage = new ChatMessage();
        chatMessage.setConversationId(conversationId);
        chatMessage.setReplyId(replyId);
        chatMessage.setUserId(userId);
        chatMessage.setRoleId(roleId);
        chatMessage.setType(messageType.getValue());
        chatMessage.setModel(model.getModel());
        chatMessage.setModelId(model.getId());
        chatMessage.setContent(content);
        chatMessage.setUseContext(useContext);
        chatMessageMapper.insert(chatMessage);
        return chatMessage;
    }
}
